import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import os
import time
from utils.utils import *
from model.DCdetector import DCdetector
from model.ChannelSim import * 
from model.ChannelTransformer import ChannelTransformer2D 
from data_factory.data_loader import get_loader_segment
from einops import rearrange
from sklearn.metrics import precision_recall_fscore_support

from metrics.metrics import *
import warnings
warnings.filterwarnings('ignore')
class PointAdjustment:
    def __init__(self,window_size=90) -> None:
        self.window_size = window_size
        self.res = [] 
        self.pred = [] 

    
    def update(self,pred,true):
        right_batch = (pred == true).astype(int)
        for i in range(right_batch.shape[0]):
            if right_batch[i] == 0: # wrong prediction
                self.res.extend([0]*self.window_size)
            else:
                self.res.extend([1]*self.window_size)

    def to_label(self,labels):
        labels = list(map(int,labels.tolist()))
        size = len(labels)
        self.res.extend([1] * (size - len(self.res)))
        for label,corr in zip(labels,self.res):
            if corr ==1 :
                self.pred.append(label)
            else:
                self.pred.append(int(not label))
        
    def point_adjustment(self,labels,f=True):
        pred = np.array(self.pred.copy()).astype(int)
        gt = labels.astype(int)
        if f :
            anomaly_state = False
            for i in range(len(gt)):
                if gt[i] == 1 and pred[i] == 1 and not anomaly_state:
                    anomaly_state = True
                    for j in range(i, 0, -1):
                        if gt[j] == 0:
                            break
                        else:
                            if pred[j] == 0:
                                pred[j] = 1
                    for j in range(i, len(gt)):
                        if gt[j] == 0:
                            break
                        else:
                            if pred[j] == 0:
                                pred[j] = 1
                elif gt[i] == 0:
                    anomaly_state = False
                if anomaly_state:
                    pred[i] = 1
        accuracy = accuracy_score(gt, pred)
        precision, recall, f_score, support = precision_recall_fscore_support(gt, pred, average='binary')
        return precision,recall,f_score,accuracy
def my_kl_loss(p, q):
    res = p * (torch.log(p + 0.0001) - torch.log(q + 0.0001))
    return torch.mean(torch.sum(res, dim=-1), dim=1)

def adjust_learning_rate(optimizer, epoch, lr_):
    lr_adjust = {epoch: lr_ * (0.5 ** ((epoch - 1) // 1))}
    if epoch in lr_adjust.keys():
        lr = lr_adjust[epoch]
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

class EarlyStopping:
    def __init__(self, patience=7, verbose=False, dataset_name='', delta=0):
        self.patience = patience
        self.verbose = verbose
        self.counter = 0
        self.best_score = None
        self.best_score2 = None
        self.early_stop = False
        self.val_loss_min = np.Inf
        self.val_loss2_min = np.Inf
        self.delta = delta
        self.dataset = dataset_name

    def __call__(self, val_loss, val_loss2, model, path):
        score = -val_loss
        score2 = -val_loss2
        if self.best_score is None:
            self.best_score = score
            self.best_score2 = score2
            self.save_checkpoint(val_loss, val_loss2, model, path)
        elif score < self.best_score + self.delta or score2 < self.best_score2 + self.delta:
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_score = score
            self.best_score2 = score2
            self.save_checkpoint(val_loss, val_loss2, model, path)
            self.counter = 0

    def save_checkpoint(self, val_loss, val_loss2, model, path):
        torch.save(model.state_dict(), os.path.join(path, str(self.dataset) + '_checkpoint.pth'))
        self.val_loss_min = val_loss
        self.val_loss2_min = val_loss2

        
class Solver(object):
    DEFAULTS = {}

    def __init__(self, config):

        self.__dict__.update(Solver.DEFAULTS, **config)

        self.train_loader = get_loader_segment(self.index, 'dataset/'+self.data_path, batch_size=self.batch_size, win_size=self.win_size, mode='train', dataset=self.dataset, )
        self.vali_loader = get_loader_segment(self.index, 'dataset/'+self.data_path, batch_size=self.batch_size, win_size=self.win_size, mode='val', dataset=self.dataset)
        self.test_loader = get_loader_segment(self.index, 'dataset/'+self.data_path, batch_size=self.batch_size, win_size=self.win_size, mode='test', dataset=self.dataset)
        self.thre_loader = get_loader_segment(self.index, 'dataset/'+self.data_path, batch_size=self.batch_size, win_size=self.win_size, mode='thre', dataset=self.dataset)

        self.build_model()
        
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        
        if self.loss_fuc == 'MAE':
            self.criterion = nn.L1Loss()
        elif self.loss_fuc == 'MSE':
            self.criterion = nn.MSELoss()
        self.criterion = FeatureDistance()
        

    def build_model(self):
        # self.model = DCdetector(win_size=self.win_size, enc_in=self.input_c, c_out=self.output_c, n_heads=self.n_heads, d_model=self.d_model, e_layers=self.e_layers, patch_size=self.patch_size, channel=self.input_c)
        # self.model = ChannelTransformer(
        #     num_variates=self.input_c,
        #     lookback_len=self.win_size,
        #     depth=6,
        #     dim=1024,
        #     num_tokens_per_variate=1,
        #     dim_head=1024 // 4,
        #     heads=4,
        #     attn_dropout=0.2,
        #     ff_mult=4,
        #     ff_dropout=0.2,
        #     use_reversible_instance_norm=True,
        #     flash_attn=True
        # )
        self.model = ChannelTransformer2D(
            num_variates=self.input_c,
            lookback_len=self.win_size,
            num_time_tokens=self.win_size//10,
            depth=6,
            dim=512,
            dim_head=512 // 4,
            heads=4,
            attn_dropout=0.2,
            ff_mult=4,
            ff_dropout=0.2,
            use_reversible_instance_norm=True,
            flash_attn=True
        )
        if torch.cuda.is_available():
            self.model.cuda()
            
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
        
        
    def vali(self, vali_loader):
        self.model.eval()
        loss_1 = []
        loss_2 = []
        for i, (input_data, _) in enumerate(vali_loader):
            input_data =input_data.to(self.device)
            channel_present = self.model(input_data)
            loss = self.criterion(channel_present).mean()
            loss_1.append(loss.item())

        return np.average(loss_1), np.average(loss_2)


    def train(self):

        time_now = time.time()
        path = self.model_save_path
        if not os.path.exists(path):
            os.makedirs(path)
        early_stopping = EarlyStopping(patience=5, verbose=True, dataset_name=self.data_path)
        train_steps = len(self.train_loader)

        for epoch in range(self.num_epochs):
            iter_count = 0

            epoch_time = time.time()
            self.model.train()
            for i, (input_data, labels) in enumerate(self.train_loader):
                # batch,win_size,channel
                iter_count +=1
                input_data = input_data.to(self.device)
                out_feature = self.model(input_data)
                loss = self.criterion(out_feature).sum()

                self.optimizer.zero_grad()
                if (i + 1) % 100 == 0:
                    speed = (time.time() - time_now) / iter_count
                    left_time = speed * ((self.num_epochs - epoch) * train_steps - i)
                    print(f'\tspeed: {speed:.4f}s/iter; left time: {left_time:.4f}s')
                    iter_count = 0
                    time_now = time.time()
                loss.backward()
                self.optimizer.step()

            vali_loss1, vali_loss2 = self.vali(self.test_loader)
            print(
                "Epoch: {0}, Cost time: {1:.3f}s ".format(
                    epoch + 1, time.time() - epoch_time))
            early_stopping(vali_loss1, vali_loss2, self.model, path)
            if early_stopping.early_stop:
                break
            adjust_learning_rate(self.optimizer, epoch + 1, self.lr)

            
    def test(self):
        self.model.load_state_dict(
            torch.load(
                os.path.join(str(self.model_save_path), str(self.data_path) + '_checkpoint.pth')))
        self.model.eval()
        temperature = 50

        # (1) stastic on the train set
        attens_energy = []
        for i, (input_data, labels) in enumerate(self.train_loader):
            input_data =input_data.to(self.device)
            channel_present = self.model(input_data)
            region_consistency = self.criterion(channel_present)
            attens_energy.append(region_consistency.detach().cpu().numpy())
        print("Train set size :", len(attens_energy))
        attens_energy = np.concatenate(attens_energy, axis=0).reshape(-1)
        train_energy = np.array(attens_energy)

        # (2) find the threshold
        attens_energy = []
        for i, (input_data, labels) in enumerate(self.test_loader):
            input_data =input_data.to(self.device)
            channel_present = self.model(input_data)
            region_consistency = self.criterion(channel_present)
            attens_energy.append(region_consistency.detach().cpu().numpy())
        

        attens_energy = np.concatenate(attens_energy, axis=0).reshape(-1)
        test_energy = np.array(attens_energy)
        combined_energy = np.concatenate([train_energy, test_energy], axis=0)
        thresh = np.percentile(combined_energy, 100 - self.anormly_ratio)
        print("Threshold :", thresh)

        # (3) evaluation on the test set
        test_labels = []
        attens_energy = []
        point_labels = []
        for i, (input_data, labels) in enumerate(self.thre_loader):
            input_data =input_data.to(self.device)
            channel_present = self.model(input_data)
            region_consistency = self.criterion(channel_present)
            attens_energy.append(region_consistency.detach().cpu().numpy())
            test_labels.append(np.clip(labels.sum(dim=1),0,1))
            point_labels.append(labels.detach().cpu().numpy())
            
        attens_energy = np.concatenate(attens_energy, axis=0).reshape(-1)
        test_labels = np.concatenate(test_labels, axis=0).reshape(-1)
        point_labels = np.concatenate(point_labels,axis=0).reshape(-1)

        point_labels = np.array(point_labels).astype(int)
        test_energy = np.array(attens_energy)
        test_labels = np.array(test_labels)

        pred = (test_energy > thresh).astype(int)
        gt = test_labels.astype(int)
        result = PointAdjustment(self.win_size)
        result.update(pred,gt)
        result.to_label(point_labels)
        precision,recall,f_score,accuracy = result.point_adjustment(point_labels)
        print("Accuracy : {:0.4f}, Precision : {:0.4f}, Recall : {:0.4f}, F-score : {:0.4f} ".format(accuracy, precision, recall, f_score))
        
        # matrix = [self.index]
        # scores_simple = combine_all_evaluation_scores(pred, gt, test_energy)
        # for key, value in scores_simple.items():
        #     matrix.append(value)
        #     print('{0:21} : {1:0.4f}'.format(key, value))

        # anomaly_state = False
        # for i in range(len(gt)):
        #     if gt[i] == 1 and pred[i] == 1 and not anomaly_state:
        #         anomaly_state = True
        #         for j in range(i, 0, -1):
        #             if gt[j] == 0:
        #                 break
        #             else:
        #                 if pred[j] == 0:
        #                     pred[j] = 1
        #         for j in range(i, len(gt)):
        #             if gt[j] == 0:
        #                 break
        #             else:
        #                 if pred[j] == 0:
        #                     pred[j] = 1
        #     elif gt[i] == 0:
        #         anomaly_state = False
        #     if anomaly_state:
        #         pred[i] = 1

        pred = np.array(pred)
        gt = np.array(gt)

        from sklearn.metrics import precision_recall_fscore_support
        from sklearn.metrics import accuracy_score

        accuracy = accuracy_score(gt, pred)
        TN = np.sum((pred == 0) & (gt == 0))
        FP = np.sum((pred == 1) & (gt == 0))
        FN = np.sum((pred == 0) & (gt == 1))
        TP = np.sum((pred == 1) & (gt == 1))
        print("TP: {}, TN: {}, FP: {}, FN: {}".format(TP, TN, FP, FN))
        precision, recall, f_score, support = precision_recall_fscore_support(gt, pred, average='binary')
        print("Accuracy : {:0.4f}, Precision : {:0.4f}, Recall : {:0.4f}, F-score : {:0.4f} ".format(accuracy, precision, recall, f_score))
        
        # if self.data_path == 'UCR' or 'UCR_AUG':
        #     import csv
        #     with open('result/'+self.data_path+'.csv', 'a+') as f:
        #         writer = csv.writer(f)
        #         writer.writerow(matrix)

        return accuracy, precision, recall, f_score
